from elasticsearch import Elasticsearch, AsyncElasticsearch, helpers

import config
import settings


class ElasticsearchStorage:
    def __init__(self):
        self.client = Elasticsearch([{'host': settings.ES_HOST, 'port': settings.ES_PORT, 'scheme': "http"}],
                                    http_auth=(settings.ES_USER, settings.ES_PASSWORD),
                                    timeout=3600)

        self.async_client = AsyncElasticsearch(
            [{'host': settings.ES_HOST, 'port': settings.ES_PORT, 'scheme': "http"}],
            http_auth=(settings.ES_USER, settings.ES_PASSWORD),
            timeout=3600)

    def bulk(self, actions=None):
        """
        创建索引
        """
        helpers.bulk(self.client, actions=actions)

    def create_index(self, index_name, index_mapping=None):
        """
        创建索引
        """
        self.client.indices.create(index=index_name, body=index_mapping)

    def delete_index(self, index_name):
        """
        删除索引
        """
        self.client.indices.delete(index=index_name)

    def index_document(self, index_name, document):
        if not settings.ONLINE:
            return
        """
        索引文档
        """
        self.client.index(index=index_name, body=document)

    def search_documents(self, index_name, query):
        if not settings.ONLINE:
            return
        """
        查询文档
        """
        response = self.client.search(index=index_name, body=query)
        return response['hits']['hits']

    def search_documents_by_agg(self, index_name, aggquery):
        """
        查询文档
        """
        response = self.client.search(index=index_name, body=aggquery)
        return response["aggregations"]["sessions"]["buckets"]

    def update_document(self, index_name, document_id, updated_fields):
        """
        更新文档
        """
        self.client.update(index=index_name, id=document_id, body={"doc": updated_fields})

    def delete_document(self, index_name, document_id):
        """
        删除文档
        """
        self.client.delete(index=index_name, id=document_id)

    async def async_create_index(self, index_name, index_mapping=None):
        """
        创建索引
        """
        await self.async_client.indices.create(index=index_name, body=index_mapping)

    async def async_delete_index(self, index_name):
        """
        删除索引
        """
        await self.async_client.indices.delete(index=index_name)

    async def async_index_document(self, index_name, document):
        if not settings.ONLINE:
            return
        """
        索引文档
        """
        await self.async_client.index(index=index_name, document=document)

    async def async_search_documents(self, index_name, query):
        if not settings.ONLINE:
            return
        """
        查询文档
        """
        response = await  self.async_client.search(index=index_name, body=query)
        return response['hits']['hits']

    async def async_search_documents_by_agg(self, index_name, aggquery):
        """
        查询文档
        """
        response = await self.async_client.search(index=index_name, body=aggquery)
        return response["aggregations"]["sessions"]["buckets"]

    async def async_update_document(self, index_name, document_id, updated_fields):
        """
        更新文档
        """
        await self.async_client.update(index=index_name, id=document_id, body={"doc": updated_fields})

    async def async_delete_document(self, index_name, document_id):
        """
        删除文档
        """
        await self.async_client.delete(index=index_name, id=document_id)


if __name__ == '__main__':
    es = ElasticsearchStorage()

    index_name = 'conversation_qw'
    index_mapping = {
        "mappings": {
            "properties": {
                "session_id": {"type": "keyword"},
                "message_id": {"type": "keyword"},
                "role": {"type": "keyword"},
                "msg_type": {"type": "keyword"},
                "attitude": {"type": "text"},
                "text_info": {"type": "text", "fields": {
                    "keyword": {
                        "type": "keyword"
                    }
                }
                              },
                "message_time": {"type": "date"},
                "customer_id": {"type": "text"}
            }
        }
    }

    # es.delete_index(index_name)

    # 创建索引
    es.create_index(index_name, index_mapping)

    # 索引文档
    # document = {
    #     'title': 'Document Title',
    #     'content': 'This is the content of the document.'
    # }
    # es.index_document(index_name, document)
    # 查询文档
    # query = {
    #     'query': {
    #         'match': {
    #             'text_info': '你好',
    #             "session_id": "JBaF4Z1Thg2NWDJSvHjo"
    #         }
    #     }
    # }
    # results = es.search_documents(index_name, query)
    # for result in results:
    #     print(result['_source'])
    #
    # # 更新文档
    # document_id = results[0]['_id']
    # updated_fields = {
    #     'attitude': '2'
    # }
    # es.update_document(index_name, document_id, updated_fields)

    # 删除文档
    # es.delete_document(index_name, '')

    # 删除索引
    # es.delete_index(index_name)

# index_name = 'conversation_qw'
# index_mapping = {
#     "mappings": {
#         "properties": {
#             "session_id": {"type": "keyword"},
#             "message_id": {"type": "keyword"},
#             "role": {"type": "keyword"},
#             "msg_type": {"type": "keyword"},
#             "attitude": {"type": "text"},
#             "text_info": {"type": "text", "fields": {
#                 "keyword": {
#                     "type": "keyword"
#                 }
#             }
#                           },
#             "message_time": {"type": "date"},
#             "customer_id":{"type": "keyword"}
#         }
#     }
# }

# index_name = 'dialogues'
# index_mapping = {
#     "mappings": {
#         "properties": {
#             "session_id": {"type": "keyword"},
#             "message_id": {"type": "keyword"},
#             "role": {"type": "text"},
#             "msg_type": {"type": "text"},
#             "attitude": {"type": "text"},
#             "text_info": {"type": "text", "fields": {
#                 "keyword": {
#                     "type": "keyword"
#                 }
#             }
#                           },
#             "message_time": {"type": "date"}
#         }
#     }
# }

# index_name = 'user_sessions'
# index_mapping = {
#     "mappings": {
#         "properties": {
#             "user_name": {"type": "keyword"},
#             "session_id": {"type": "keyword"},
#             "session_title": {"type": "text"},
#             "text_info": {"type": "text", "fields": {
#                 "keyword": {
#                     "type": "keyword"
#                 }
#             }
#                           },
#             "message_time": {"type": "date"}
#         }
#     }
# }
